from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import os

# 初始化DeepSeek AI模型
llm = ChatOpenAI(
    model="deepseek-chat",
    openai_api_key=os.getenv("OPENAI_API_KEY_DEEPSEEK"),
    openai_api_base="https://api.deepseek.com/v1"
)

# 创建聊天提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个AI助手"),
    ("human", "{input}")
])

# 创建聊天链
chain = prompt | llm


def main():
    print("DeepSeek聊天助手已启动，输入'quit'退出")
    while True:
        user_input = input("你: ")
        if user_input.lower() == 'quit':
            print("聊天结束")
            break

        response = chain.invoke({"input": user_input})
        print(f"AI助手: {response.content}")


if __name__ == '__main__':
    main()
